#!/usr/bin/env python3
# set coding: utf-8
"""
@Project : openresty
@File :    litewafserver.py
@Time :    2022/12/20 9:29
@Author :  richard zhu
@Email :   gaotao.zhugt@dtzhejiang.com
@purpose : 
"""
from flask import Flask, request, jsonify
import json
app = Flask(__name__)

@app.route('/api/heartbeat', methods=['POST'])
def heartbeat():
    # print(request.values)
    # print(data.get('api_key'))
    data = request.json
    api_key = data['api_key']
    api_password = data['api_password']
    if api_key and api_password == 'd13f3851-efa7-4c33-b6f7-c3f5e1f7d701':
        return jsonify(data='normal', result=True)
    else:
        return jsonify(data='error', result=False)

@app.route('/api/config/global', methods=['POST'])
def conf_global():
    """
    全局配置是需要litewaf重启的
    :return:
    """
    data = request.json
    api_key = data['api_key']
    api_password = data['api_password']
    if api_key and api_password == 'd13f3851-efa7-4c33-b6f7-c3f5e1f7d701':
        # 域名对应的规则类型
        rs = {
                "default_page": {"code": "404", "html":"Not in default ServerName"},
                "redis": {"host":"127.0.0.1", "port":6379, "password":""},
                "sys_abnormal_handle_data": {"bypass_check":True, "same_name_args_check":True, "truncated_agrs_check":True,
                                             "client_body_size_check":True, "code": "401", "html":"litewaf request abnormal handle"},
                "waf_module_enable" : ["web_white_check","web_black_check","domain_check","flow_protection_check","flow_limit_burst_check"],
                "auto_update_period": 300,
                "flow_protection": {"default_limit": 1000, "code":"401", "html":'<html><head><meta http-equiv="Content-Type" content="text/html; charset=utf-8" /><meta http-equiv="Content-Language" content="zh-cn" /><title>网站防火墙</title></head><body><h1 align="center"> 网站waf防火墙已拦截 </body></html>'}
                # 默认每个域名1w/s的QPS，可以通过规则再调整
        }
        return jsonify(configure_without_change=False,
                       global_rule_data=rs,
                       result=True)
    else:
        return jsonify(data='auth error', result=False)

@app.route('/api/config/domain', methods=['POST'])
def conf_domain():
    data = request.json
    api_key = data['api_key']
    api_password = data['api_password']
    if api_key and api_password == 'd13f3851-efa7-4c33-b6f7-c3f5e1f7d701':
        # 域名对应的规则类型
        d2t = {
                "ip.gaotao.club": {"rules": {"flow_limit_burst_rule": ["1"], "flow_protection_rule": []}},
                "wp.gaotao.club": {"rules": {"flow_limit_burst_rule": ["2"]}},
                "*.test.com": {"rules": {}}   # 没有数据也要这样回
        }

        # [{"match_args": [{"key": "http_args", "value": "query_string"}], "args_prepocess": ["base64Decode"],"match_operator": "rx", "match_value": "select.+(from|limit)"}]
        # [{"match_args":[{"key":"http_args","value":"src_ip"}],"args_prepocess":["none"],"match_operator":"str_eq","match_value":"127.0.0.1"}]
        r = {
            "web_white_rule": {"1": {"rule_name": "ip_white_rule1", "rule_matchs": [
                {"match_args": [{"key": "http_args", "value": "src_ip"}], "args_prepocess": ["none"],
                 "match_operator": "ipmask", "match_value": "127.0.0.2/24"}]}}, # ipmask必须是一个列表

            "flow_limit_burst_rule": {"1": {"rule_name": "limit_ip.gaotao.club_redis",
                                            "action_value": [{"rate": 2, "burst": 1}],
                                            "rule_action": "brust_limit_redis",
                                            "rule_matchs": [
                                                    {"match_args": [{"key": "http_args", "value": "host"}], "args_prepocess": ["none"],
                                                     "match_operator": "str_eq", "match_value": "ip.gaotao.club"}]},
                                      
                                      "2": {"rule_name": "limit_wp.gaotao.club_local",
                                            "action_value": [{"rate": 2, "burst": 1}],
                                            "rule_action": "brust_limit_local",
                                            "rule_matchs": [
                                                    {"match_args": [{"key": "http_args", "value": "host"}], "args_prepocess": ["none"],
                                                     "match_operator": "str_eq", "match_value": "wp.gaotao.club"}]}
                                      },
            
            "flow_protection_rule" : {"1": {"rule_name": "high_freq_cc_rate_ip.gaotao.club",
                                               "action_value": [{"limit":"180","time":"3600",
                                                                 "factor_keys": [{"key": "http_args", "value": "host"},
                                                                                 {"key": "http_args", "value": "src_ip"}]
                                                                 }],
                                               "rule_action": "brust_limit_redis",  # block 或者 人机监测 或者放通
                                               "rule_matchs": [
                                                   {"match_args": [{"key": "http_args", "value": "host"}],
                                                    "args_prepocess": ["none"],
                                                    "match_operator": "str_eq", "match_value": "ip.gaotao.club"}
                                               ]

                                            }
                                      }
        }


        return jsonify(configure_without_change=False,
                       domain_rule_data=d2t,
                       rules_data = r,
                       auto_update_period=30,
                       result=True)
    else:
        return jsonify(data='auth error', result=False)



if __name__ == '__main__':
    app.run(host='0.0.0.0', port=1000, debug=True)
